from __future__ import annotations

import json
import os
from collections import defaultdict

import torch

from generation.asset_converter import AssetConverterFactory
from generation.enums import AssetType, LayoutInfo, Scene3DItemEnum
from metasim.constants import PhysicStateType
from metasim.scenario.objects import RigidObjCfg

__all__ = [
    "cvt_embodiedgen_asset_to_anysim",
    "load_embodiedgen_asset",
    "load_embodiedgen_layout_pose",
]


def cvt_embodiedgen_asset_to_anysim(
    urdf_files: list[str],
    target_dirs: list[str],
    target_type: AssetType,
    source_type: AssetType,
    overwrite: bool = False,
    **kwargs,
) -> dict[str, str]:
    """Convert URDF files generated by EmbodiedGen into the format required by all simulators.

    Supported simulators include SAPIEN, Isaac Sim, MuJoCo, Isaac Gym, Genesis, and Pybullet.

    Example:
        dst_asset_path = cvt_embodiedgen_asset_to_anysim(
            urdf_files,
            target_dirs,
            target_type=SimAssetMapper[simulator_name],
            source_type=AssetType.MESH,
        )

    Args:
        urdf_files (List[str]): List of URDF file paths to be converted.
        target_dirs (List[str]): List of target directories for converted assets.
        target_type (AssetType): The target asset type.
        source_type (AssetType): The source asset type.
        overwrite (bool): Whether to overwrite existing converted files.
        **kwargs: Additional keyword arguments for the converter.

    Returns:
        Dict[str, str]: A dictionary mapping the original URDF file path to the converted asset file path.
    """
    if isinstance(urdf_files, str):
        urdf_files = [urdf_files]
    if isinstance(target_dirs, str):
        urdf_files = [target_dirs]

    # If the target type is URDF, no conversion is needed.
    if target_type == AssetType.URDF:
        return {key: key for key in urdf_files}

    asset_converter = AssetConverterFactory.create(
        target_type=target_type,
        source_type=source_type,
        **kwargs,
    )
    asset_paths = dict()

    with asset_converter:
        for urdf_file, target_dir in zip(urdf_files, target_dirs):
            filename = os.path.basename(urdf_file).replace(".urdf", "")
            if target_type == AssetType.MJCF:
                target_file = f"{target_dir}/{filename}.xml"
            elif target_type == AssetType.USD:
                target_file = f"{target_dir}/{filename}.usd"
            else:
                raise NotImplementedError(f"Target type {target_type} not supported.")
            if not os.path.exists(target_file) or overwrite:
                asset_converter.convert(urdf_file, target_file)

            asset_paths[urdf_file] = target_file

    return asset_paths


def load_embodiedgen_asset(
    layout: str | LayoutInfo,
    target_type: AssetType,
    source_type: AssetType,
    **kwargs,
) -> list[RigidObjCfg]:
    """Load assets from an EmbodiedGen layout and convert them to the specified format."""
    asset_root = os.path.dirname(layout)
    if isinstance(layout, str) and layout.endswith(".json"):
        layout = LayoutInfo.from_dict(json.load(open(layout)))

    asset_info = defaultdict(dict)
    for node in layout.assets:
        file_dir = layout.assets[node]
        file_name = f"{node.replace(' ', '_')}.urdf"
        if layout.objs_mapping[node] == Scene3DItemEnum.BACKGROUND.value:
            continue

        kinematic_static = layout.relation.get(Scene3DItemEnum.CONTEXT.value, None) == node

        urdf_file = os.path.join(asset_root, file_dir, file_name)
        asset_info[urdf_file]["kinematic_static"] = kinematic_static
        node_name = node.replace(" ", "_")
        asset_info[urdf_file]["name"] = node_name

    urdf_files = list(asset_info.keys())
    target_dirs = [f"{os.path.dirname(p)}/../{target_type}" for p in urdf_files]
    asset_paths = cvt_embodiedgen_asset_to_anysim(urdf_files, target_dirs, target_type, source_type, **kwargs)
    rigid_cfgs = []
    for urdf_file in asset_paths:
        rigid_cfg = RigidObjCfg(
            name=asset_info[urdf_file]["name"],
            scale=(1, 1, 1),
            physics=PhysicStateType.RIGIDBODY,
            fix_base_link=asset_info[urdf_file]["kinematic_static"],
        )
        if target_type == AssetType.MJCF:
            asset_key = "mjcf_path"
        elif target_type == AssetType.USD:
            asset_key = "usd_path"
        elif target_type == AssetType.URDF:
            asset_key = "urdf_path"
        setattr(rigid_cfg, asset_key, asset_paths[urdf_file])
        rigid_cfgs.append(rigid_cfg)

    return rigid_cfgs


def load_embodiedgen_layout_pose(
    layout: str | LayoutInfo,
    z_offset: float = 0.0,
    robot_delta_z: float = 0.01,
) -> dict:
    """Load the initial poses of objects and robots from an EmbodiedGen layout."""
    if isinstance(layout, str) and layout.endswith(".json"):
        layout = LayoutInfo.from_dict(json.load(open(layout)))

    init_states = [{"objects": {}, "robots": {}}]
    for node in layout.assets:
        if layout.objs_mapping[node] == Scene3DItemEnum.BACKGROUND.value:
            continue

        # Objects initial states
        x, y, z, qx, qy, qz, qw = layout.position[node].copy()
        if layout.objs_mapping[node] != Scene3DItemEnum.CONTEXT.value:
            z += z_offset

        node_name = node.replace(" ", "_")
        init_states[0]["objects"][node_name] = {
            "pos": torch.tensor([x, y, z]),
            "rot": torch.tensor([qw, qx, qy, qz]),
        }

        # Robot initial states
        robot_name = layout.relation[Scene3DItemEnum.ROBOT.value]
        x, y, z, qx, qy, qz, qw = layout.position[robot_name]
        init_states[0]["robots"]["default"] = {
            "pos": torch.tensor([x, y, z + robot_delta_z]),
            "rot": torch.tensor([qw, qx, qy, qz]),
        }

    return init_states
